'''
https://leetcode.cn/problems/maximum-average-pass-ratio/
'''
import heapq
from typing import List

class Heap:
    def __init__(self, nums=None, comparator=lambda x, y: x < y):
        self.nums = list(nums) if nums else []
        self.size = len(self.nums)
        self.comparator = comparator
        for i in range((self.size - 2) // 2, -1, -1):
            self._heapify(i)

    def get_size(self):
        return self.size

    def empty(self):
        return self.size == 0

    def peek(self):
        if self.size == 0:
            return None
        return self.nums[0]

    def pop(self):
        if self.size == 0:
            return None
        nums = self.nums
        nums[0], nums[-1] = nums[-1], nums[0]
        res = nums.pop()
        self.size -= 1
        if self.size > 0:
            self._heapify(0)
        return res

    def push(self, v):
        self.nums.append(v)
        self._heap_insert(self.size)
        self.size += 1

    def pop_push(self, v):
        if self.empty():
            self.push(v)
            return None
        res = self.nums[0]
        self.nums[0] = v
        self._heapify(0)
        return res

    def push_pop(self, v):
        if self.empty() or self.comparator(v, self.nums[0]):
            return v
        res = self.nums[0]
        self.nums[0] = v
        self._heapify(0)
        return res

    def _heap_insert(self, i):
        nums = self.nums
        while (parent := (i - 1) // 2) >= 0 and self.comparator(nums[i], nums[parent]):
            nums[i], nums[parent] = nums[parent], nums[i]
            i = parent

    def _heapify(self, i):
        nums = self.nums
        size = self.size
        while (l := i * 2 + 1) < size:
            r = l + 1
            best = r if (r < size and self.comparator(nums[r], nums[l])) else l
            best = best if self.comparator(nums[best], nums[i]) else i
            if best == i:
                break
            nums[i], nums[best] = nums[best], nums[i]
            i = best

class Solution:
    def maxAverageRatio(self, classes: List[List[int]], extraStudents: int) -> float:
        # 1: 通过人数，2:总人数， 0: 再来一个人的通过率提升了多少
        heap = Heap(comparator=lambda x, y: x[0] > y[0])
        for n_pass, total in classes:
            inc_ratio = (n_pass + 1) / (total + 1) - n_pass / total
            heap.push((inc_ratio, n_pass, total))
        for i in range(extraStudents):
            inc_ratio, n_pass, total = heap.pop()
            n_pass, total = n_pass + 1, total + 1
            inc_ratio = (n_pass + 1) / (total + 1) - n_pass / total
            heap.push((inc_ratio, n_pass, total))
        return sum(temp[1] / temp[2] for temp in heap.nums) / heap.get_size()

    # python 堆
    def maxAverageRatio2(self, classes: List[List[int]], extraStudents: int) -> float:
        # 1: 通过人数，2:总人数， 0: 再来一个人的通过率提升了多少
        heap = [(
            -((n_pass + 1) / (total + 1) - n_pass / total), n_pass, total
        ) for n_pass, total in classes]
        heapq.heapify(heap)

        for i in range(extraStudents):
            _, n_pass, total = heapq.heappop(heap)
            n_pass, total = n_pass + 1, total + 1
            heapq.heappush(heap, (-((n_pass + 1) / (total + 1) - n_pass / total), n_pass, total))
        return sum(n_pass / total for _, n_pass, total in heap) / len(heap)

classes = [[1,2],[3,5],[2,2]]
extraStudents = 2
print(Solution().maxAverageRatio(classes, extraStudents))
print(Solution().maxAverageRatio2(classes, extraStudents))